from llvmlite import ir

from .base import BuiltinFunction, get_builtin_function_cls
from MelodieAllegro.types import integer, double, string


class Malloc(BuiltinFunction):
    def register(self):
        func_ty = ir.FunctionType(ir.PointerType(ir.IntType(8)), [integer])
        self.func = ir.Function(self.module, func_ty, name="pm_mem_allocate")


class Realloc(BuiltinFunction):
    def register(self):
        func_ty = ir.FunctionType(ir.PointerType(ir.IntType(8)), [ir.PointerType(ir.IntType(8)), integer, integer])
        self.func = ir.Function(self.module, func_ty, name="pm_mem_reallocate")


class AllocAutoPointer(BuiltinFunction):
    def register(self):
        func_ty = ir.FunctionType(ir.PointerType(ir.IntType(8)), [integer])
        self.func = ir.Function(self.module, func_ty, name="pm_auto_ptr_allocate")


class PMRuntimeInit(BuiltinFunction):
    def register(self):
        func_ty = ir.FunctionType(integer, [integer])
        self.func = ir.Function(self.module, func_ty, name="pm_runtime_init")


class RuntimeInit(BuiltinFunction):
    def register(self):
        from ..builtin_func_types import builtins
        func_ty = ir.FunctionType(integer, [integer])
        self.func = ir.Function(self.module, func_ty, name="runtime_init")
        bb = self.func.append_basic_block('entry')
        builder = ir.IRBuilder(bb)

        builder.call(builtins.pm_runtime_init.func, [integer(0)])
        builder.ret(integer(0))


class PrintString(BuiltinFunction):
    def register(self):
        func_ty = ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.IntType(8))])
        self.func = ir.Function(self.module, func_ty, name="print_string")

    def __call__(self, builder: ir.IRBuilder, content, name: str = ""):
        assert content.type == string
        i8_ptr = ir.IntType(8).as_pointer()
        ptr2 = builder.bitcast(content, i8_ptr)
        builder.call(self.func, [content])


class PrintLF(BuiltinFunction):
    def register(self):
        func_ty = ir.FunctionType(ir.VoidType(), [])
        self.func = ir.Function(self.module, func_ty, name="print_lf")

    def __call__(self, builder: ir.IRBuilder):
        builder.call(self.func, [])


class PrintDouble(BuiltinFunction):
    def register(self):
        func_ty = ir.FunctionType(ir.VoidType(), [double])
        self.func = ir.Function(self.module, func_ty, name="print_double")


class PrintInteger(BuiltinFunction):
    def register(self):
        func_ty = ir.FunctionType(ir.VoidType(), [integer])
        self.func = ir.Function(self.module, func_ty, name="print_integer")


class RaiseIndexError(BuiltinFunction):
    def register(self):
        func_ty = ir.FunctionType(ir.VoidType(), [integer, integer, integer])
        self.func = ir.Function(self.module, func_ty, name="pm_raise_index_error")

    def __call__(self, builder: ir.IRBuilder, low, high, index):
        builder.call(self.func, [low, high, index])


class RaiseError(BuiltinFunction):
    def register(self):
        func_ty = ir.FunctionType(ir.VoidType(), [ir.PointerType(ir.IntType(8))])
        self.func = ir.Function(self.module, func_ty, name="raise_error")

    def __call__(self, builder: ir.IRBuilder, content: str, name: str = ""):
        typ = ir.ArrayType(ir.IntType(8), len(content))
        val = ir.Constant(typ,
                          bytearray(content.encode("utf8")))
        ptr = builder.alloca(typ)
        builder.store(val, ptr)

        i8_ptr = ir.IntType(8).as_pointer()
        ptr2 = builder.bitcast(ptr, i8_ptr)

        builder.call(self.func, [
            ptr2
        ])


builtin_runtime_funcs = get_builtin_function_cls(globals())
